import os
import sys
import time
import uuid
import shutil
import subprocess

import click
import yaml
from loguru import logger
from multiprocessing import Pool, Manager

from module_analysis import get_top_modules



def RED_INFO(info, **karg):
    print(f"\033[0;31m{info}\033[0m", **karg)

def GREEN_INFO(info, **karg):
    print(f"\033[0;32m{info}\033[0m", **karg)

def YELLOW_INFO(info, **karg):
    print(f"\033[0;33m{info}\033[0m", **karg)

def BLUE_INFO(info, **karg):
    print(f"\033[0;34m{info}\033[0m", **karg)


class YosysCase:
    def __init__(self, case_path, platform):
        self.case_path = case_path
        self.platform = platform
        self.init()

    def init(self):
        self.name = os.path.basename(self.case_path)
        self.orfs_dir = os.path.join(self.case_path, self.platform,'yosys', 'orfs')
        self.orfs_dc_dir = os.path.join(self.case_path, self.platform,'yosys', 'orfs-dc')

        self.verilog_dir = os.path.join(self.case_path, 'src')
        self.top_module_name = get_top_modules(self.verilog_dir)[0]

        self.uuid_name = str(uuid.uuid4())

    def create_orfs_dir(self):
        if not os.path.exists(self.orfs_dir):
            os.makedirs(self.orfs_dir)

    def create_orfs_dc_dir(self):
        if not os.path.exists(self.orfs_dc_dir):
            os.makedirs(self.orfs_dc_dir)

    @staticmethod
    def get_cases(case_path_list, case_platform):
        """Preprocess the path of the cases and convert them to YosysCase objects."""
        cases = []
        for case_path in case_path_list:
            cases.append(YosysCase(case_path, case_platform))

        return cases

    def generate_configmk(self, configmk_path):
        lines = [
            f"export DESIGN_NICKNAME = {self.uuid_name}\n",
            f"export DESIGN_NAME = {self.top_module_name}\n",
            f"export PLATFORM = {self.platform}\n",
            "\n",
            "export VERILOG_FILES = $(sort $(wildcard ./designs/src/$(DESIGN_NICKNAME)/*.v))\n",
            "export SDC_FILE = ./designs/$(PLATFORM)/$(DESIGN_NICKNAME)/constraint.sdc\n",
        ]
        with open(configmk_path, "w") as file:
            file.writelines(lines)

    def generate_sdc(self, sdc_path):
        lines = []
        with open(sdc_path, "w") as file:
            file.writelines(lines)

    def genrate_orfs_run_log(self, run_text):
        run_log = os.path.join(self.orfs_dir, 'run.log')
        with open(run_log, "w", encoding="utf-8") as file:
            file.write(run_text)

    def process_configmk(self, configmk_path):
        with open(configmk_path, 'r') as file:
            lines = file.readlines()

        for i, line in enumerate(lines):
            if line.startswith('export DESIGN_NICKNAME'):
                lines[i] = f"export DESIGN_NICKNAME = {self.name}\n"

        with open(configmk_path, 'w') as file:
            file.writelines(lines)

    def generate_orfsdc_run_log(self, run_text):
        run_log = os.path.join(self.orfs_dc_dir, 'run.log')
        with open(run_log, "w", encoding="utf-8") as file:
            file.write(run_text)


class Result:
    PASS = 0
    FAILED = 1
    INVALID = 2

    def __init__(self, case_name):
        self.name = case_name
        self.info = []
        self.state = None


class OrfsDcRunner:
    def __init__(self, orfs_flow_dir):
        self.orfs_flow_dir = orfs_flow_dir
        self.name = "run orfs get area report and run dc get timing report"

    def generate_sdc(self):
        pass
    
    def generate_dc_runtcl(self, synth_verilog_path, case: YosysCase):
        dc_runtcl_path = os.path.join(case.orfs_dc_dir, 'run.tcl')
        log_path = os.path.join(case.orfs_dc_dir, 'log')
        os.mkdir(log_path)
        
        lines = [
            "set target_library NangateOpenCellLibrary_typical.db\n",
            "set link_library [concat $target_library \"*\"]\n",
            "set search_path [concat $search_path /home/share/platform/nangate45/lib]\n",
            "define_design_lib WORK -path work_lib\n",
            f"read_file -format verilog {synth_verilog_path}\n",
            f"analyze -format sverilog {synth_verilog_path}\n",
            f"elaborate {case.top_module_name}\n",
            f"current_design {case.top_module_name}\n",
            "read_sdc ../constraint.sdc\n"
            "\n",
            "report_area > log/area.rpt\n",
            "report_design > log/design.rpt\n",
            "report_timing -path end > log/timing.rpt\n",
            "report_timing -max_path 10 >> log/timing.rpt\n",
            "report_power > log/power.rpt\n",
            "report_qor > log/qor.rpt\n",
            "exit\n"
        ]

        with open(dc_runtcl_path, 'w') as file:
            file.writelines(lines)

    def clean_dc_tempfile(self, case: YosysCase):
        command_log = os.path.join(case.orfs_dc_dir, 'command.log')
        work_lib = os.path.join(case.orfs_dc_dir, 'work_lib')
        default_svf = os.path.join(case.orfs_dc_dir, 'default.svf')

        if os.path.exists(command_log):
            os.remove(command_log)

        if os.path.exists(default_svf):
            os.remove(default_svf)

        if os.path.exists(work_lib):
            shutil.rmtree(work_lib)


    def run(self, case: YosysCase):
        result = OrfsRunner(self.orfs_flow_dir).run(case)
        result.state = None
        result.info.clear()

        case.create_orfs_dc_dir()

        synth_verilog_path = os.path.join(case.orfs_dir, 'result', 'base', '1_synth.v')
        if not os.path.exists(synth_verilog_path):
            result.info.append(f"[x_x] FAILED : {case.name}")
            result.info.append(f"FAILED REASON : openroad flow script not generate netlist")
            result.state = Result.INVALID
            return result
        
        sdc_path = os.path.join(case.case_path, case.platform,'constraint.sdc')
        if not os.path.exists(sdc_path):
            self.generate_sdc()

        self.generate_dc_runtcl(synth_verilog_path, case)
        
        ret = subprocess.run(['dc_shell', '-f', 'run.tcl'], cwd=case.orfs_dc_dir, capture_output=True, text=True)
        case.generate_orfsdc_run_log(ret.stdout + ret.stderr)

        self.clean_dc_tempfile(case)

        result.state = Result.PASS
        result.info.append(f"[^_^] Orfs-Dc RUN PASS : {case.name}")
        result.info.append(f"orfs run result directory : {case.orfs_dir}")
        result.info.append(f"dc run result directory : {case.orfs_dc_dir}")
        return result


class OrfsRunner:
    def __init__(self, orfs_flow_dir):
        self.orfs_flow_dir = orfs_flow_dir
        self.name = "run orfs"
        
    def run(self, case: YosysCase):
        case.create_orfs_dir()

        orfs_verilog_dir = os.path.join(self.orfs_flow_dir, 'designs', 'src', case.uuid_name)
        shutil.copytree(case.verilog_dir, orfs_verilog_dir)

        orfs_config_dir = os.path.join(self.orfs_flow_dir, 'designs', case.platform, case.uuid_name)
        os.mkdir(orfs_config_dir)
        configmk_path = os.path.join(orfs_config_dir, 'config.mk')
        case.generate_configmk(configmk_path)
        case.generate_sdc(os.path.join(orfs_config_dir, 'constraint.sdc'))

        configmk_relative_path = './designs/' + configmk_path.split('/designs/')[1]
        ret = subprocess.run(['make', f'DESIGN_CONFIG={configmk_relative_path}'], cwd=self.orfs_flow_dir, capture_output=True, text=True)
        case.genrate_orfs_run_log(ret.stdout + ret.stderr)
        
        log_path = os.path.join(self.orfs_flow_dir, 'logs', case.platform, case.uuid_name)
        object_path = os.path.join(self.orfs_flow_dir, 'objects', case.platform, case.uuid_name)
        report_path = os.path.join(self.orfs_flow_dir, 'reports', case.platform, case.uuid_name)
        result_path = os.path.join(self.orfs_flow_dir, 'results', case.platform, case.uuid_name)

        shutil.copytree(log_path, os.path.join(case.orfs_dir, 'log'))
        shutil.copytree(object_path, os.path.join(case.orfs_dir, 'object'))
        shutil.copytree(report_path, os.path.join(case.orfs_dir, 'report'))
        shutil.copytree(result_path, os.path.join(case.orfs_dir, 'result'))
        
        case.process_configmk(configmk_path)
        shutil.copytree(orfs_config_dir, os.path.join(case.orfs_dir, 'config'))

        shutil.rmtree(log_path)
        shutil.rmtree(object_path)
        shutil.rmtree(report_path)
        shutil.rmtree(result_path)
        shutil.rmtree(orfs_verilog_dir)
        shutil.rmtree(orfs_config_dir)

        result = Result(case.name)
        result.state = Result.PASS
        result.info.append(f"[^_^] Orfs RUN PASS : {case.name}")
        result.info.append(f"run result directory : {case.orfs_dir}")
        return result
        

class YosysTestFrame:
    def __init__(self):
        self.cpu_num = None
        self.platform = None
        self.orfs_flow_dir = None
        self.flow = None

        self.case_type = YosysCase
        
        self.cmd_width = os.get_terminal_size().columns

    def set_flow(self, flow):
        self.flow = flow

    def set_platform(self, platform):
        self.platform = platform
    
    def set_cpu_num(self, cpu_num):
        self.cpu_num = cpu_num

    def set_orfs_flow_dir(self, orfs_flow_dir):
        self.orfs_flow_dir = orfs_flow_dir

    def _pre_run_cases(self):
        if os.path.exists(self.orfs_flow_dir):
            BLUE_INFO("-" * self.cmd_width)
            logger.info(f"SET PDK: {self.platform}")
            BLUE_INFO("-" * self.cmd_width)
        else:
            logger.error("openRoad flow script not exists")
            sys.exit(1)

    def _init_process(self):
        pass

    def async_run_case(self, case: YosysCase, queue):
        try:
            result = self.flow(self.orfs_flow_dir).run(case)

        except Exception as e:
            logger.error(f"setup case <{case.name}> fail: {e}")
            result = Result(case.name)
            result.info.append(f"[x_x] SETUP FAILED : {case.name}")
            result.info.append(f"FAILED REASON      : {e}")
            result.state = Result.INVALID
        except KeyboardInterrupt as e:
            logger.error(f'user stopped the case "{case.name}"')
            result = Result(case.name)
            result.info.append(f"[x_x] RUN STOPPED   : {case.name}")
            result.state = Result.INVALID

        try:
            queue.put(result)
        except Exception as e:
            logger.error(f"put result to queue error: {e}")

    def process_result(self, queue):
        pass_rate = None
        pass_case_count = 0
        fail_case_count = 0
        invalid_case_count = 0
        fail_case_set = set()
        invalid_case_set = set()
        error_cases = []

        if queue.empty():
            logger.error("NOT FOUND ANY CASE, PLEASE CHECK CASE NAME.")

        while not queue.empty():
            result = queue.get()

            if result.state == Result.INVALID:
                invalid_case_count += 1
                invalid_case_set.add(result.name)
                error_cases.append(result)
            elif result.state == Result.FAILED:
                fail_case_count += 1
                fail_case_set.add(result.name)
                error_cases.append(result)
            else:
                pass_case_count += 1
                GREEN_INFO("-" * self.cmd_width)
                print("\n".join(result.info))
                GREEN_INFO("-" * self.cmd_width)

        for error_case in error_cases:
            RED_INFO("-" * self.cmd_width)
            print("\n".join(error_case.info))
            RED_INFO("-" * self.cmd_width)

        # 统计通过率
        try:
            pass_rate = pass_case_count / (
                pass_case_count + fail_case_count + invalid_case_count
            )
        except ZeroDivisionError:
            pass_rate = 0

        # 进行结果输出和记录
        output_list = []

        if pass_rate < 1:
            output_list.append(f"PASS RATE  : {pass_rate * 100:.2f}%")
            output_list.append(f"PASSED     : {pass_case_count}")
            output_list.append(f"FAILED     : {fail_case_count}")
            output_list.append(f"INVALID    : {invalid_case_count}")
            output_list.append(f"")
            output_list.append(f"FAILED  CASE: " + " ".join(fail_case_set))
            output_list.append(f"INVALID CASE: " + " ".join(invalid_case_set))

            RED_INFO("#" * self.cmd_width)
            RED_INFO("\n".join(output_list))
            RED_INFO("#" * self.cmd_width)
        else:
            output_list.append(f"PASS RATE  : {pass_rate * 100:.2f}%")
            output_list.append(f"PASSED     : {pass_case_count}")
            output_list.append(f"FAILED     : {fail_case_count}")
            output_list.append(f"INVALID    : {invalid_case_count}")
            output_list.append(f"")
            output_list.append(f"FAILED  CASE: " + " ".join(fail_case_set))
            output_list.append(f"INVALID CASE: " + " ".join(invalid_case_set))

            GREEN_INFO("#" * self.cmd_width)
            GREEN_INFO("\n".join(output_list))
            GREEN_INFO("#" * self.cmd_width)

        return fail_case_count + invalid_case_count

    def async_run_cases(self, all_cases: list):
        cases = self.case_type.get_cases(all_cases, self.platform)

        self._pre_run_cases()

        pool = Pool(
            processes=self.cpu_num,
            initializer=self._init_process,
        )
        result_queue = Manager().Queue()

        for case in cases:
            pool.apply_async(self.async_run_case, args=(case, result_queue))

        pool.close()
        stop_en = False
        while not stop_en:
            try:
                pool.join()
                stop_en = True
            except KeyboardInterrupt:
                try:
                    last_stop = time.time()
                    time.sleep(0.5)
                    pool.join()
                except KeyboardInterrupt:
                    stop_en = True

        error_case_num = self.process_result(result_queue)
        return error_case_num
    
    def run_cases(self, cases: list):
        try:
            return self.async_run_cases(cases)
        except Exception as e:
            logger.error(f"run case error: {e}")
            return 1
        


@click.command()
@click.argument("cases", type=click.STRING, nargs=-1)  # Absolute path of cases
@click.option("--cpu_num", type=click.INT, default = 4, help = "Specified the count of processes")
@click.option("--pdk", type=click.Choice(["asap7", "sky130hd", "nangate45"]), default="nangate45", help='Specified PDK_name')
def main(
    cases, 
    cpu_num, 
    pdk,
):
    tf = YosysTestFrame()

    tf.set_cpu_num(cpu_num)
    tf.set_platform(pdk)
    tf.set_orfs_flow_dir("/home/share/process/OpenROAD-flow-scripts/flow")
    tf.set_flow(OrfsDcRunner)  # OrfsRunner OrfsDcRunner

    sys.exit(tf.run_cases(list(cases)))


if __name__ == "__main__":
    main()    
